import sys
import os
import cv2
import numpy as np
from PIL import Image
from typing import Union
import open3d as o3d
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import argparse
from tqdm import tqdm

def read_rgb(rgb_path):
    rgb_image = Image.open(rgb_path).convert('RGB')
    return np.array(rgb_image)

def read_depth(depth_path):
    depth_image = Image.open(depth_path)
    return np.array(depth_image)

def display_point_cloud(depth_path, guide_path, focal_lenght, depth_unit, display_title='', depth_display_range=None, export_dir=None):
    depth_image = read_depth(depth_path)
    guide_color_image = None if guide_path is None else read_rgb(guide_path)
    # 读取深度图像
    if depth_image is None:
        print(f"无法读取深度")
        return

    if depth_display_range is None:
        depth_display_range = [depth_image.min(), depth_image.max()]

    # 显示范围内的深度值归一化
    min_display_depth, max_display_depth = depth_display_range
    depth_image = np.clip(depth_image, min_display_depth, max_display_depth)
    # 获取图像的尺寸
    height, width = depth_image.shape
    # 生成点云
    points = []
    colors = []
    fx, fy = focal_lenght, focal_lenght  # 焦距，可以根据实际相机参数调整
    cx, cy = width / 2, height / 2  # 主点，可以根据实际相机参数调整

    # 获取深度值的最大值和最小值，用于归一化颜色映射
    min_depth = np.min(depth_image[depth_image > 0])
    max_depth = np.max(depth_image)
    scale = 1
    if depth_unit == 'mm':
        scale = 0.001
    elif depth_unit == 'm':
        scale = 1
    for v in range(height):
        for u in range(width):
            z = depth_image[v, u] * scale  # 将深度值从毫米转换为米
            if z <= 0:  # 跳过无效的深度值
                continue
            x = (u - cx) * z / fx
            y = (v - cy) * z / fy
            points.append([x, y, z])
            if guide_color_image is not None:
                color = guide_color_image[v,u]
                color = color / 255
            else:
                # 归一化深度值，范围在0到1之间
                normalized_depth = (depth_image[v, u] - min_depth) / (max_depth - min_depth)
                # 将归一化深度值映射到颜色（蓝色到红色的渐变）
                color = plt.cm.jet(normalized_depth)[:3]  # 使用jet颜色映射，取前3个值作为RGB
            colors.append(color)

    # 创建Open3D点云对象
    point_cloud = o3d.geometry.PointCloud()
    point_cloud.points = o3d.utility.Vector3dVector(np.array(points))
    point_cloud.colors = o3d.utility.Vector3dVector(np.array(colors))
    if export_dir is not None:
        name = os.path.splitext(os.path.basename(depth_path))[0]
        output_path = os.path.join(export_dir, name + '.ply')
        o3d.io.write_point_cloud(output_path, point_cloud)
    else:
        # 创建坐标轴
        coordinate_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=1.0, origin=[0, 0, 0])
        # 显示点云和坐标轴
        o3d.visualization.draw_geometries([point_cloud, coordinate_frame], window_name = display_title +' - '+ depth_path, 
                                        width=800, height=600, left=50, top=50, 
                                        point_show_normal=False, mesh_show_wireframe=False, mesh_show_back_face=False)

def display_depth_map(depth_path, guide_path, depth_unit, display_title='', depth_display_range=None, export_dir=None):
    guide_color_image = None if guide_path is None else read_rgb(guide_path)
    depth_image = read_depth(depth_path)
    # 读取深度图像
    if depth_image is None:
        print(f"无法读取深度")
        return
    
    if depth_display_range is None:
        depth_display_range = [depth_image.min(), depth_image.max()]

    # 显示范围内的深度值归一化
    min_display_depth, max_display_depth = depth_display_range
    depth_image = np.clip(depth_image, min_display_depth, max_display_depth)
    depth_image_normalized = (depth_image - min_display_depth) / (max_display_depth - min_display_depth)

    # 使用HSV颜色空间将深度值映射到彩色
    colormap = plt.get_cmap('jet')
    depth_colored = colormap(depth_image_normalized)[:, :, :3]  # 忽略alpha通道

    # 创建颜色条
    norm = mcolors.Normalize(vmin=min_display_depth, vmax=max_display_depth)
    sm = plt.cm.ScalarMappable(cmap=colormap, norm=norm)
    sm.set_array([])

    # 使用gridspec创建布局
    fig = plt.figure(figsize=(12, 6))
    gs = fig.add_gridspec(1, 3, width_ratios=[1, 1, 0.05], wspace=0.3)

    # 左侧显示guide_color_image
    if guide_color_image is not None:
        ax0 = fig.add_subplot(gs[0, 0])
        ax0.imshow(guide_color_image)
        ax0.set_title('Color Image')
        ax0.axis('off')

    # 右侧显示depth_image
    ax1 = fig.add_subplot(gs[0, 1])
    ax1.imshow(depth_colored)
    ax1.set_title(display_title + ' - ' + depth_path)
    ax1.axis('off')

    # 在右侧的一个小区域添加颜色条
    cbar_ax = fig.add_subplot(gs[0, 2])
    cbar = plt.colorbar(sm, cax=cbar_ax)
    cbar.set_label('Depth Value')

    if export_dir is not None:
        name = os.path.splitext(os.path.basename(depth_path))[0]
        output_path = os.path.join(export_dir, name + '.png')
        plt.savefig(output_path)
    else:
        plt.show()

def main():
    # 创建 ArgumentParser 对象
    parser = argparse.ArgumentParser(description="mask show")
    parser.add_argument("filepath", help="文件的路径")
    parser.add_argument("-g", "--guide_color_image", type=str, help="set guide color image")
    parser.add_argument("-f", "--focal_lenght", type=float, default=200, help="set default focal lenght")
    parser.add_argument("-c", "--display_point_cloud", action='store_true', help="set if display depth point cloud points")
    parser.add_argument("-r", "--depth_display_range", nargs='+', type=int, default=None, help="set display depth value range")
    parser.add_argument("-e", "--export_dir", type=str, default=None, help="set export dir")
    parser.add_argument("-u", "--depth_unit", type=str, default='mm', help="set depth value unit")
    parser.add_argument("-t", "--display_title", type=str, default='', help="set window title")
    parser.add_argument("-s", "--suffixes", nargs='+', default=['jpg', 'jpeg', 'png', 'tif', 'pgm', 'ppm'], type=str, help="image file suffixes")
    args = parser.parse_args()
    depth_paths = []
    guide_paths:Union[None|dict|str] = None
    if os.path.isdir(args.filepath):
        depth_paths = [os.path.join(args.filepath, f) for f in os.listdir(args.filepath) if f.endswith(tuple(args.suffixes))]
    elif args.filepath.endswith(tuple(args.suffixes)):
        depth_paths.append(args.filepath)

    if args.guide_color_image is not None:
        if os.path.isdir(args.guide_color_image):
            guide_paths = {}
            for f in os.listdir(args.guide_color_image):
                if f.endswith(tuple(args.suffixes)):
                    name, _ = os.path.splitext(os.path.basename(f))
                    guide_paths[name] = os.path.join(args.guide_color_image, f)
        elif args.guide_color_image.endswith(tuple(args.suffixes)):
            guide_paths = args.guide_color_image

    for i, depth_path in enumerate(tqdm(depth_paths)):
        if guide_paths is None:
            guide_path = None
        elif isinstance(guide_paths, str):
            guide_path = guide_paths
        elif isinstance(guide_paths, dict):
            name,_ = os.path.splitext(os.path.basename(depth_path))
            guide_path = None if name not in guide_paths else guide_paths[name]
        if args.display_point_cloud:
            display_point_cloud(depth_path, guide_path, focal_lenght=args.focal_lenght, depth_unit=args.depth_unit, display_title = args.display_title, depth_display_range=args.depth_display_range, export_dir=args.export_dir)
        else:
            display_depth_map(depth_path, guide_path, depth_unit=args.depth_unit, display_title = args.display_title, depth_display_range=args.depth_display_range, export_dir=args.export_dir)

# 使用示例
if __name__ == "__main__":
    main()